import datetime
import functools
import logging
import signal
from absl import app, flags
from tornado.ioloop import PeriodicCallback

from coin.base.timestamp import get_timestamp
from coin.base.param_util import to_list
from coin.exchange.base.order_gateway_config import convert_og_config_from_dict
from coin.exchange.base.strategy_logger import StrategyLogProvider
from coin.exchange.kr_rest.product.product_impl import (generate_product_from_subreq)
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.exchange.huobi_futures.kr_rest.futures_product import HuobiFuturesProduct
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.util.feed_checker import get_default_feed_checker_config
from coin.proto.coin_feed_pb2 import BestBookEntry, BookEntry
from coin.proto.coin_request_pb2 import AccountRequestProto
from coin.proto.coin_telemetry_pb2 import LatencyProto

from coin.strategy.accounting.pnl_balance_printer.util import (get_quote_currency, reserve_as_proto)
from coin.strategy.executor.passive.util import (convert_executor_config_from_dict)
from coin.strategy.executor.passive.factory import construct_mm_executor
from coin.strategy.executor.util.cancel_all_executor import (SimpleCancelAllExecutor)
from coin.strategy.hk_hard_arb.index_pricer import IndexPricer
from coin.strategy.hk_hard_arb.pricer import PositionLeanPricer
from coin.strategy.hk_hard_arb.util import load_clean_json
from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.mm.base.unary_multifeed_strategy_base import UnaryMultifeedStrategyBase

from coin.strategy.mm.subscription import (
    FeedSubscriptionRequest,
    SubscriptionRequest,
)
from coin.strategy.tool.fill_log_dumper import FillLogDumper
from coin.strategy.hk_hard_arb.amplitude_checker import AmplitudeChecker
from coin.strategy.accounting.pnl_balance_printer.factory import get_pnl_balance_printer_cls
from coin.strategy.util.request_util import convert_strategy_request_from_dict
from coin.support.telemetry.latency_recorder import LatencyRecorder, LifeOfSignal

FLAGS = flags.FLAGS


def _within_rollover_period(product, timestamp):
  ret = False
  dt = datetime.datetime.utcfromtimestamp(timestamp / 10**9)
  if isinstance(product, (OkexFuturesProduct, HuobiFuturesProduct)):
    if dt.isoweekday() == 5:
      if (dt.hour == 7 and dt.minute >= 55) or \
         (dt.hour == 8 and dt.minute <= 5):
        ret = True
  return ret


def _check_strat_config(strat_config):
  assert 'common' in strat_config
  assert 'products' in strat_config
  assert 'trade_subreq' in strat_config['common']
  assert 'strategy_request' in strat_config['common']
  assert isinstance(strat_config['products'], dict)
  assert len(strat_config['products']) > 0


def _check_product_config(config):
  assert 'trade_subreq' in config
  assert 'trade_symbol' in config
  assert 'ref_feed_subreqs' in config
  assert 'ref_symbol' in config
  assert len(config['ref_symbol']) == len(config['ref_feed_subreqs'])
  assert 'basis_ma_window_min' in config


def _adjust_price_by_bound(price, upper_bound, lower_bound):
  if upper_bound is not None and lower_bound is not None:
    assert upper_bound > lower_bound, (upper_bound, lower_bound)
  if price is not None:
    if upper_bound is not None:
      price = min(price, upper_bound)
    if lower_bound is not None:
      price = max(price, lower_bound)
  return price


class PassUnhedgeProfile(object):
  def __init__(self, config_dict, start_time=None):
    _check_product_config(config_dict)
    self.trade_subreq = SubscriptionRequest.from_dict(config_dict['trade_subreq'])
    self.trade_product = generate_product_from_subreq(self.trade_subreq,
                                                      config_dict['trade_symbol'],
                                                      config_dict.get('trading_date', start_time))
    self.currency = list({self.trade_product.base, get_quote_currency(self.trade_product)})
    ref_subreq = []
    ref_product = []
    for subreq_dict, symbol in zip(config_dict['ref_feed_subreqs'], config_dict['ref_symbol']):
      subreq = SubscriptionRequest.from_dict(subreq_dict)
      product = generate_product_from_subreq(subreq, symbol)
      ref_subreq.append(subreq)
      ref_product.append(product)
    self.ref_subreq = ref_subreq
    self.ref_product = ref_product
    self.config_dict = config_dict


class PassUnhedgeStrategyConfig(object):
  def __init__(self, config, start_time=None):
    _check_strat_config(config)
    self._config = config
    self._trade_subreq = SubscriptionRequest.from_dict(config['common']['trade_subreq'])
    self._model_profile = self._gen_model_profile(start_time)
    self._feed_checker_config = self._gen_feed_checker_config()
    self._strategy_request = convert_strategy_request_from_dict(
        config['common']['strategy_request'])

  @property
  def config(self):
    return self._config

  @property
  def strategy_request(self):
    return self._strategy_request

  @property
  def trade_subreq(self):
    return self._trade_subreq

  @property
  def feed_checker_config(self):
    return self._feed_checker_config

  @property
  def og_params(self):
    og_params = self._config['common'].get('og_params', {})
    og_params = convert_og_config_from_dict(self._trade_subreq, og_params)
    return og_params

  @property
  def model_profile(self):
    return self._model_profile

  def _gen_model_profile(self, start_time=None):
    profile_dict = {}
    common_config = self._config['common']
    for symbol, config in self._config['products'].items():
      trade_symbol = config['trade_symbol']
      product = generate_product_from_subreq(self._trade_subreq, trade_symbol, start_time)
      config = {**common_config, **config}
      assert symbol not in profile_dict
      profile = PassUnhedgeProfile(config, start_time)
      assert profile.trade_product.full_symbol == product.full_symbol
      profile_dict[symbol] = profile
    return profile_dict

  def _gen_feed_checker_config(self):
    check_ref_feed = self._config['common'].get('ref_exchange_feed_checker', False)
    check_trade_feed = self._config['common'].get('trade_exchange_feed_checker', False)
    feed_checker_products = []
    for profile in self.model_profile.values():
      if check_ref_feed:
        feed_checker_products += to_list(profile.ref_product)
      if check_trade_feed:
        feed_checker_products += to_list(profile.trade_product)
    feed_checker_config = get_default_feed_checker_config(feed_checker_products)
    return feed_checker_config

  def get_products(self):
    return [profile.trade_product for profile in self._model_profile.values()]

  def get_currencies(self):
    currencies = []
    for profile in self._model_profile.values():
      currencies.extend(profile.currency)
    return list(set(currencies))


class SingleProductSubStrategy(object):
  def __init__(self, *, model_profile, dry_run, logger, strategy_logger, latency_recorder, parent):
    self._dry_run = dry_run
    self._logger = logger
    self._strategy_logger = strategy_logger
    self._latency_recorder = latency_recorder
    self._parent = parent
    config = model_profile.config_dict
    self._config = config
    self._trade_subreq = model_profile.trade_subreq
    self._trade_product = model_profile.trade_product
    self._ref_subreq = model_profile.ref_subreq
    self._ref_product = model_profile.ref_product
    assert len(self._ref_subreq) == len(self._ref_product)

    self._edge = None
    self._spread = None
    self._spread_hard = None
    self._pull_edge = None

    self._bbo_protection = config.get('bbo_protection', False)
    if 'edge_bp' in config:
      self._edge = config['edge_bp'] / 10000.
    if 'spread_bp' in config:
      self._spread = config['spread_bp'] / 10000.
    if 'spread_bp_hard' in config:
      self._spread_hard = config['spread_bp_hard'] / 10000.
    if 'pull_edge_bp' in config:
      self._pull_edge = config['pull_edge_bp'] / 10000.
    if 'pull_spread_bp' in config:
      self._pull_edge = config['pull_spread_bp'] / 10000. / 2

    if 'fill_widen_within_sec' in config:
      self._fill_widen_within = config['fill_widen_within_sec'] * 1e9
      self._fill_widen = config['fill_widen_bp'] / 10000.
    else:
      self._fill_widen_within = None
      self._fill_widen = None

    self._amplitude_checker = None
    if 'amplitude_checker' in config:
      checker_config = config['amplitude_checker']
      window_min = checker_config['window_min']
      threshold_bp = checker_config['threshold_bp']
      self._amplitude_checker = AmplitudeChecker(self._trade_product, window_min, threshold_bp)

    self._close_only_with_position_threshold = config.get('close_only_with_position_threshold',
                                                          None)
    if self._close_only_with_position_threshold is not None:
      assert self._close_only_with_position_threshold >= 0
    self._within_rollover_period = False

    assert 1 == sum([
        self._edge is not None,
        self._spread is not None,
    ])

    products = self._ref_product + to_list(self._trade_product)
    self._bookmap = BookMap(products, logger=self._logger)

    tick = get_holder_from_product(self._trade_product).product_info.gen_tick()
    self._pricer = IndexPricer(self._config, tick)
    self._pos_lean_pricer = PositionLeanPricer(min_pos=config['min_pos'],
                                               max_pos=config['max_pos'],
                                               lot_size=config['lot_size'],
                                               lean_edge_bp=config.get('lean_edge_bp', 0))

    self._last_book_write_ts = get_timestamp()

    self._executor = None
    self._exit = False

  def _check_rollover_period(self):
    within_rollover_period = _within_rollover_period(self._trade_product, get_timestamp())
    if within_rollover_period:
      self._logger.info('Within rollover period.')
    if self._within_rollover_period and not within_rollover_period:
      self._parent.register_exit()
      self._logger.info('Exit after rollover')
    self._within_rollover_period = within_rollover_period

  def _close_only_enforced(self):
    amplitude_high = False
    if self._amplitude_checker is not None:
      amplitude_high = self._amplitude_checker.alarm
    close_only_with_position = False
    if self._close_only_with_position_threshold is not None:
      pos = self._executor.query_unbiased_position() - \
            self._executor.config2.reserve
      if abs(pos) > self._close_only_with_position_threshold:
        close_only_with_position = True
    close_only = \
        self._within_rollover_period or amplitude_high or \
        close_only_with_position
    return close_only

  def set_exit(self):
    self._exit = True

  def on_og_reset(self, *, order_gateway):
    self._logger.info('Intializing executor...')
    executor_params = convert_executor_config_from_dict(self._trade_product.symbol, self._config)
    self._executor = construct_mm_executor(self._trade_subreq,
                                           executor_params,
                                           order_gateway,
                                           self._logger,
                                           latency_recorder=self._latency_recorder)
    self._strategy_logger.write_executor_config(executor_params)

  def prepare(self, feed_subsystem):
    assert feed_subsystem is not None
    product = self._trade_product
    subreq = self._trade_subreq
    trade_feed_sub_request = FeedSubscriptionRequest.create(to_list(product), subreq)
    feed_subsystem.subscribe_book(trade_feed_sub_request,
                                  functools.partial(self.on_book, product, self._trade_subreq))
    for product, sub_req in zip(self._ref_product, self._ref_subreq):
      feed_sub_request = FeedSubscriptionRequest.create(to_list(product), sub_req)
      feed_subsystem.subscribe_book(feed_sub_request,
                                    functools.partial(self.on_book, product, sub_req))
    if self._amplitude_checker is not None:
      feed_subsystem.subscribe_book(
          trade_feed_sub_request,
          functools.partial(self._amplitude_checker.on_book, self._trade_product))
      PeriodicCallback(self._amplitude_checker.update_alarm, 10 * 1000).start()
    PeriodicCallback(self._check_rollover_period, 10 * 1000).start()

  def on_book(self, product, sub_req, book):
    if self._exit:
      return

    if product == self._trade_product and self._executor is not None:
      book_ts = book.timestamp
      if book_ts - self._last_book_write_ts >= 5 * 10**9:
        self._last_book_write_ts = book_ts
        best_book_entry = BestBookEntry(symbol=product.symbol,
                                        fetched_ts=book_ts,
                                        ask0=BookEntry(price_float=book.ask0().price),
                                        bid0=BookEntry(price_float=book.bid0().price))
        self._strategy_logger.write_best_book_entry(
            best_book_entry, account_request=self._executor.og.og_logger.request)

    if self._latency_recorder is not None:
      account_request = AccountRequestProto(market_type=sub_req.market_type,
                                            exchange=sub_req.exchange,
                                            api_version=sub_req.api_version)
      self._latency_recorder.begin(LatencyProto.FEED_LATENCY)
      self._latency_recorder.set_time(LifeOfSignal.RAW_FEED_PARSE_STARTED,
                                      account_request=account_request,
                                      timestamp=book.timestamp)
      self._latency_recorder.set_time(
          LifeOfSignal.ON_BOOK_FEED,
          account_request=account_request,
      )
      self._latency_recorder.end()
      self._latency_recorder.begin()
      self._latency_recorder.set_time(LifeOfSignal.RAW_FEED_PARSE_STARTED,
                                      account_request=account_request,
                                      timestamp=book.timestamp)
      self._latency_recorder.set_time(
          LifeOfSignal.ON_BOOK_FEED,
          account_request=account_request,
      )

    self._bookmap.set_book(product, book)

    # Check book freshness and sanity
    if not self._bookmap.is_book_ready() or not self._bookmap.is_book_sane():
      return

    ref_book = [self._bookmap.get_book(product) for product in self._ref_product]
    trade_book = self._bookmap.get_book(self._trade_product)
    self._pricer.update(ref_book, trade_book)

    if not self._pricer.ready or self._executor is None:
      return

    if self._spread is not None:
      askp, bidp = self._pricer.calculate_price_by_spread(self._spread,
                                                          bbo_protection=self._bbo_protection)
    elif self._edge is not None:
      askp, bidp = self._pricer.calculate_price(self._edge,
                                                self._edge,
                                                bbo_protection=self._bbo_protection)
    else:
      raise ValueError('not supported mode')

    close_only = self._close_only_enforced()
    self._executor.set_close_only(close_only)
    if close_only:
      askp = book.ask0().price
      bidp = book.bid0().price
    else:
      pos = self._executor.query_unbiased_position()
      askp, bidp = self._pos_lean_pricer.update(pos=pos, buy_price=bidp, sell_price=askp)

    # self._executor._lfillt has some issue in compare position to get last fill time.
    # that will cause last fill time always to be latest time and fill widen is always
    # enabled. so this functionality is muted.
    """
    if self._fill_widen_within and \
        get_timestamp() - self._executor._lfillt._ts._last_timestamp < \
            self._fill_widen_within:
      askp = askp * (1 + self._fill_widen)
      bidp = bidp * (1 - self._fill_widen)
    """

    upper_bound = None
    lower_bound = None
    if self._spread is not None and self._spread_hard is not None:
      upper_bound, lower_bound = self._pricer.calculate_price_by_spread(
          self._spread_hard, bbo_protection=self._bbo_protection)
    askp = _adjust_price_by_bound(price=askp, upper_bound=upper_bound, lower_bound=lower_bound)
    bidp = _adjust_price_by_bound(price=bidp, upper_bound=upper_bound, lower_bound=lower_bound)
    askp, bidp = self._pricer.adjust_price(askp_1=askp,
                                           bidp_1=bidp,
                                           bbo_protection=self._bbo_protection)

    sell_price_pull = None
    buy_price_pull = None
    if askp is not None and self._pull_edge is not None:
      sell_price_pull = askp * (1 - self._pull_edge)
    if bidp is not None and self._pull_edge is not None:
      buy_price_pull = bidp * (1 + self._pull_edge)

    if self._dry_run:
      self._logger.info('product: %s, sell: %s, buy: %s, sell_pull: %s, buy_pull: %s' %
                        (self._trade_product, askp, bidp, sell_price_pull, buy_price_pull))
    else:
      self._executor.manage_mm_orders(sell_price=askp,
                                      buy_price=bidp,
                                      sell_price_pull=sell_price_pull,
                                      buy_price_pull=buy_price_pull)


class PassUnhedgeStrategy(UnaryMultifeedStrategyBase):
  def __init__(self, strat_config, dry_run=True, exit_after_min=None, logger=None):
    self._logger = logger or logging.getLogger('PassUnhedgeStrategy')
    assert isinstance(strat_config, PassUnhedgeStrategyConfig)
    self._config = strat_config
    self._set_traits(self._config.trade_subreq)
    self._exit_after_min = exit_after_min
    super().__init__(products=self._config.get_products(),
                     currencies=self._config.get_currencies(),
                     logger=self._logger)

    self._book = BookMap(self._products, logger=self._logger)

    self._reserve_map = self._gen_reserve_map()
    printer_cls = get_pnl_balance_printer_cls(self._config.trade_subreq)
    self._printer = printer_cls(self._products, self._book, self._reserve_map, logger=self._logger)

    self._fill_dumper = FillLogDumper(self._config.trade_subreq)

    self._order_canceled_after_og_ready = False
    self._exit_attempt = 0
    self._exit_executor = None
    self._strategy_logger = StrategyLogProvider(self._config.strategy_request)
    self._strategy_logger.write_request()
    if strat_config.config['common'].get('enable_telemetry'):
      self._latency_recorder = LatencyRecorder('life_of_signal', self._strategy_logger)
      self._logger.info('telemetry enabled!')
    else:
      self._latency_recorder = None
      self._logger.info('telemetry disabled!')

    self._single_product_sub_strategies = []
    for profile in self._config.model_profile.values():
      self._single_product_sub_strategies.append(
          SingleProductSubStrategy(model_profile=profile,
                                   dry_run=dry_run,
                                   logger=self._logger,
                                   strategy_logger=self._strategy_logger,
                                   latency_recorder=self._latency_recorder,
                                   parent=self))

  @property
  def config(self):
    return self._config

  @property
  def strategy_logger(self):
    return self._strategy_logger

  def _gen_reserve_map(self):
    reserve_map = {}
    if self._config.trade_subreq.market_type == 'Spot':
      for profile in self._config.model_profile.values():
        currency = profile.trade_product.base.currency
        min_pos = profile.config_dict['min_pos']
        max_pos = profile.config_dict['max_pos']
        reserve = 0.5 * (min_pos + max_pos)
        reserve_map[currency] = reserve
    return reserve_map

  def dump_pnl_balance(self):
    self._printer.print_pnl_balance()
    pnl_balances = self._printer.calculate_pnl_balance()
    if pnl_balances is not None:
      for pnl_balance in pnl_balances:
        self._strategy_logger.write_pnl_balance(pnl_balance=pnl_balance, quote=None, subreq_id=None)

  def prepare(self):
    for profile in self.config.model_profile.values():
      product = profile.trade_product
      feed_sub_request = FeedSubscriptionRequest.create(to_list(product), profile.trade_subreq)
      self.feed_subsystem.subscribe_book(feed_sub_request, functools.partial(self.on_book, product))
    for strat in self._single_product_sub_strategies:
      strat.prepare(self.feed_subsystem)
    # Start dumping pnl balance every 30 seconds
    PeriodicCallback(self.dump_pnl_balance, 30 * 1000).start()
    if self._exit_after_min is not None:
      self.register_exit(datetime.timedelta(minutes=self._exit_after_min))

  def on_og_reset(self):
    if not self._order_canceled_after_og_ready:
      self.order_gateway.cancel_all()
      self._logger.info('Orders canceled after OGs are ready.')
      self._order_canceled_after_og_ready = True

    self._printer.register_order_gateway(self.order_gateway)
    for strat in self._single_product_sub_strategies:
      strat.on_og_reset(order_gateway=self.order_gateway)
    self.order_gateway.register_fill_callback(self._fill_dumper.dump_fill)

    self._exit_executor = SimpleCancelAllExecutor()

    if self._reserve_map:
      self._strategy_logger.write_reserve(reserve_as_proto(self._config.trade_subreq,
                                                           self._reserve_map),
                                          account_request=self.order_gateway.og_logger.request)

    # TODO(chensili): Add OkexFutures limit error check

  def on_book(self, product, book):
    self._book.set_book(product, book)

  def signal_handler(self, sig, frame):
    self._logger.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self.on_exit)

  def register_exit(self, timedelta=None):
    if timedelta is None:
      self._ioloop.add_callback(self.on_exit)
    else:
      self._ioloop.add_timeout(timedelta, self.on_exit)

  def on_exit(self):
    for strat in self._single_product_sub_strategies:
      strat.set_exit()

    self._exit_attempt += 1
    if self._exit_attempt > 5:
      self._ioloop.add_callback(self._ioloop.stop)

    finished = True
    if self._exit_executor is not None:
      ogs = to_list(self.order_gateway)
      self._logger.info('Exit callback executed.')
      finished = self._exit_executor.cancel_all_on_exit(ogs)

    if finished:
      self._ioloop.add_callback(self._ioloop.stop)
    else:
      timeout = datetime.timedelta(seconds=2)
      self._ioloop.add_timeout(timeout, self.on_exit)


def main(argv):
  init_og_after_min = 0.1

  strat_config = PassUnhedgeStrategyConfig(load_clean_json(FLAGS.config_filename))
  strategy = PassUnhedgeStrategy(strat_config,
                                 dry_run=FLAGS.dry_run,
                                 exit_after_min=FLAGS.exit_after_min)

  signal.signal(signal.SIGTERM, strategy.signal_handler)
  signal.signal(signal.SIGINT, strategy.signal_handler)

  feed_checker_config = strat_config.feed_checker_config
  return strategy.run_live(strategy=strategy,
                           on_order_gateway_reset_callback=strategy.on_og_reset,
                           og_config_name=FLAGS.og_config_name,
                           og_params=strategy.config.og_params,
                           init_order_gateway_after=datetime.timedelta(minutes=init_og_after_min),
                           feed_checker_config=feed_checker_config,
                           og_log_callback=strategy.strategy_logger.write_og_log,
                           slow_callback_duration=0.01)


if __name__ == '__main__':
  flags.DEFINE_string('og_config_name', None, 'og_config_name')

  flags.DEFINE_string('config_filename', '', 'config_filename')

  flags.DEFINE_integer('exit_after_min', 30, 'exit_after_min')

  flags.DEFINE_boolean('dry_run', False, 'dry_run')

  flags.DEFINE_float('slow_callback_duration', None, 'slow_callback_duration')

  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  app.run(main)
